from functools import wraps
from pathlib import Path
import pickle
from typing import Callable
import time

##---------------------------------------------------------------------------
#                            Decorators
#---------------------------------------------------------------------------

def timer(func:Callable=None, label:str=""):
    """
    Count cost time of a function and add a label. 
    
    Args:
        func (Callable): function for running
        label (str): label for print
    
    Returns:
        Callable: decorated function
    
    Example: 
    >>> @get_time
    >>> def test():
    >>>     time.sleep(1)
    >>> test()

    >>> @get_time(label="test")
    >>> def test2():
    >>>     time.sleep(1)
    >>> test2()

    >>> test3 = timer(test, label="test3")
    >>> test3()
    """
    # when time decroator is called without arguments, return a wrapper
    if func is None:
        
        return lambda f: timer(f, label=label)
    
    # when time decroator is called with arguments, return a wrapper
    def wrapper(*args, **kwargs):
        
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()

        elapsed_time = end_time - start_time
        days = int(elapsed_time // (24 * 3600))
        elapsed_time %= (24 * 3600)
        hours = int(elapsed_time // 3600)
        elapsed_time %= 3600
        minutes = int(elapsed_time // 60)
        seconds = elapsed_time % 60


        time_str = f"{days} days" if days else ""
        time_str += f"{hours} hours" if hours or days else ""
        time_str += f"{minutes} mins" if minutes or hours or days else ""
        time_str += f"{seconds:.2f} secs" if seconds or minutes or hours or days else ""

        label2 =  label if label else func.__name__
        print(f"Runing {label2}: {time_str}")
        print("\n", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        return result
    
    return wrapper

def load_pickle(save_name:str|Path = "", save=True):
    """
    用于保存和加载模型的装饰器。

    Args:
        save_name (str or Path): 模型的保存路径。
        save (bool): 是否尝试保存模型。默认为 True。

    Returns:
        装饰后的函数，具有缓存功能的函数。

    Example:
    >>> @load_pickle(save=False)
    >>> def run_model(save_name="model.pkl"):
    >>>     # Simulating some computation
    >>>     fitted = {"model": "trained"}
    >>>     return fitted
    >>> # 注意，load_pickle 不能读取默认参数值，必须手动传入 
    >>> fitted = run_model(save_name="model.pkl")
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            
            # Loading the result from a file if it exists
            # save_path = kwargs.get('save_name')
            save_path = Path(save_name)
            if save_path.exists():
                print(f"save_name is {save_path}, try to load----------------")
                try:
                    out = pickle.load(save_path.open('rb'))
                    return out
                except Exception as e:
                    raise e

            # Call the original function
            out = func(*args, **kwargs)

            # Optionally save the result
            if save and save_name != "" and out is not None:
                print(f"save_name is {save_name}, try to save----------------")
                try:
                    pickle.dump(out, save_name.open('wb'))
                except FileNotFoundError:
                    print(f"saving {save_name} failed ----------------")

            return out

        return wrapper

    return decorator


import pandas as pd
import numpy as np
from itertools import product
def generate_design_matrix(
        factor_levels, n_repeat=1, start_num=0
    ) -> np.ndarray:
        """Generate a design matrix.

        Parameters:
        factor_levels (tuple): A tuple where each element represents the number of levels for a factor.
        n_repeat (int, optional): The number of times to repeat the design matrix. Defaults to 1.
        start_num (int, optional): The starting number for factor encoding. Defaults to 0.

        Returns:
        np.ndarray: The generated design matrix.

        Examples:
        >>> _generate_design_matrix({"stim":["left","right"]})
        [[0]
        [1]]
        """

        matrix = pd.DataFrame(
            list(product(*factor_levels.values())),
            columns=list(factor_levels.keys())
        )
        matrix = matrix.apply(lambda x: pd.factorize(x)[0] + start_num)

        repeated_matrix = np.tile(matrix, (n_repeat, 1))

        return repeated_matrix